package com.iogogogo.tenant.config;

import com.google.common.collect.Maps;
import com.iogogogo.tenant.config.dao.DataSourceConfigDao;
import org.hibernate.engine.jdbc.connections.spi.AbstractDataSourceBasedMultiTenantConnectionProviderImpl;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.util.Map;

/**
 * Created by tao.zeng on 2021/8/8.
 */
@Component
public class DataSourceBasedMultiTenantConnectionProviderImpl extends AbstractDataSourceBasedMultiTenantConnectionProviderImpl implements InitializingBean {

    private final DataSource dataSource;

    private final ApplicationContext applicationContext;

    private final Map<String, DataSource> dataSourceMap = Maps.newHashMap();

    public DataSourceBasedMultiTenantConnectionProviderImpl(@Qualifier("defaultDataSource") DataSource dataSource, ApplicationContext applicationContext) {
        this.dataSource = dataSource;
        this.applicationContext = applicationContext;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        // 初始化数据源，设置默认值
        dataSourceMap.put("defaultDataSource", dataSource);
    }

    @Override
    protected DataSource selectAnyDataSource() {
        return dataSourceMap.get("defaultDataSource");
    }

    @Override
    public DataSource selectDataSource(String tenant) {
        // 根据tenant创建数据源
        DataSource dataSource = dataSourceMap.get(tenant);
        if (dataSource == null) {
            TenantDataSource tenantDataSource = applicationContext.getBean(TenantDataSource.class);
            DataSourceConfigDao dataSourceConfigDao = applicationContext.getBean(DataSourceConfigDao.class);
            dataSource = tenantDataSource.getDataSource(dataSourceConfigDao.getByName(tenant));
            dataSourceMap.put(tenant, dataSource);
        }
        return dataSource;
    }
}
